#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：yolo_train.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/1/7 上午9:22
"""

import os
import subprocess
import torch
from typing import List


class YOLOv5Trainer:
	def __init__(self, input_dir: str, output_dir: str, classes: List[str], use_gpu: bool = True):
		"""
		初始化YOLOv5训练器。

		:param input_dir: 包含训练图像和标签的目录路径。
		:param output_dir: 输出模型训练结果的目录路径。
		:param classes: 类别列表。
		:param use_gpu: 是否使用GPU，默认为True。
		"""
		self.input_dir = os.path.abspath(input_dir)
		self.output_dir = os.path.abspath(output_dir)
		self.classes = classes
		self.train_images_dir = os.path.join(self.input_dir, 'images')
		self.train_labels_dir = os.path.join(self.input_dir, 'labels')

		# 选择设备
		self.device = 'cuda' if (use_gpu and torch.cuda.is_available()) else 'cpu'

	def train_model(self):
		"""
		训练YOLOv5模型。
		"""
		data_yaml_path = os.path.join(self.output_dir, 'data.yaml')
		weights_path = 'yolov5s.pt'  # 使用预训练的YOLOv5s模型
		output_dir = os.path.join(self.output_dir, 'runs', 'train')
		img_size = 640  # 图像大小
		batch_size = 8  # 批量大小
		epochs = 5  # 训练轮数
		workers = 2  # 数据加载器的工作线程数

		# 创建data.yaml文件
		with open(data_yaml_path, 'w') as f:
			f.write(f"train: {self.train_images_dir}\n")
			f.write(f"val: {self.train_images_dir}\n")
			f.write(f"nc: {len(self.classes)}\n")
			f.write(f"names: {self.classes}\n")

		# 检查data.yaml内容
		with open(data_yaml_path, 'r') as f:
			print("data.yaml 内容:")
			print(f.read())

		# 训练命令
		yolov5_dir = os.path.dirname(os.path.abspath(__file__))  # 获取当前脚本目录
		train_command = [
			"python", os.path.join(yolov5_dir, 'yolov5', 'train.py'),
			"--img", str(img_size),
			"--batch", str(batch_size),
			"--epochs", str(epochs),
			"--data", data_yaml_path,
			"--cfg", "yolov5/models/yolov5s.yaml",
			"--weights", weights_path,
			"--device", self.device,
			"--project", output_dir,
			"--workers", str(workers)
		]
		print("训练命令:", " ".join(train_command))
		result = subprocess.run(train_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
		print("Standard Output:", result.stdout.decode())
		print("Standard Error:", result.stderr.decode())
		if result.returncode != 0:
			print("训练失败，返回码:", result.returncode)
		else:
			print("训练成功")
